// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <gtest/gtest.h>

#include <pollux/functions/lib/zeta_distribution.h>

namespace kumo::pollux::functions {
namespace {

TEST(ZetaDistributionTest, basic) {
  constexpr int N = 200;
  ZetaDistribution dist(1.1, N);
  int freq[N]{};
  unsigned seed = 0;
  // seed = std::random_device()();
  std::default_random_engine gen(seed);
  for (int i = 0; i < 100'000; ++i) {
    int k = dist(gen);
    ASSERT_GE(k, 1);
    ASSERT_LE(k, N);
    ++freq[k - 1];
  }
  // Strictly decreasing for the first 10.
  for (int i = 1; i < 10; ++i) {
    EXPECT_GT(freq[i - 1], freq[i]) << i;
  }
  // Moving window sum decreasing for the first 50.
  int sum = std::accumulate(freq, freq + 10, 0);
  for (int i = 10; i < 50; ++i) {
    int newSum = sum - freq[i - 10] + freq[i];
    EXPECT_LT(newSum, sum) << i;
    sum = newSum;
  }
  // The rest is long tail with few hits.
  for (int i = 50; i < N; ++i) {
    EXPECT_LT(freq[i], freq[40]) << i;
  }
}

} // namespace
} // namespace kumo::pollux::functions
